import asyncio
from pathlib import Path

import pandas as pd
import graphrag.api as api
from graphrag.config.load_config import load_config


async def demoGlobalSearch():
    PROJECT_DIRECTORY = "./datas/demo2"

    graphrag_config = load_config(Path(PROJECT_DIRECTORY))

    entities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/entities.parquet")
    communities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/communities.parquet")
    community_reports = pd.read_parquet(
        f"{PROJECT_DIRECTORY}/output/community_reports.parquet"
    )

    response, context = await api.global_search(
        config=graphrag_config,
        entities=entities,
        communities=communities,
        community_reports=community_reports,
        community_level=1,
        dynamic_community_selection=False,
        response_type="Multiple Paragraphs",
        query="请帮我介绍下亚马逊公司？",
    )

    print(response)

async def demoLocalSearch():
    PROJECT_DIRECTORY = "./datas/demo2"

    graphrag_config = load_config(Path(PROJECT_DIRECTORY))

    entities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/entities.parquet")
    communities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/communities.parquet")
    community_reports = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/community_reports.parquet")

    text_unit = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/text_units.parquet")
    relations_ships = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/relationships.parquet")

    response, context = await api.local_search(
        config=graphrag_config,
        entities=entities,
        communities=communities,
        community_reports=community_reports,
        community_level=2,
        text_units=text_unit,
        covariates=None,
        relationships=relations_ships,
        response_type="Multiple Paragraphs",
        #query="苹果公司推出 iPhone 的时间？",
        query="你知道DeepSeek模型吗？",
    )

    print(response)

if __name__ == "__main__":
    #asyncio.run(demoGlobalSearch())
    asyncio.run(demoLocalSearch())
